import os
import sys
sys.path.append('/share/duli/utils')
sys.path.append('/share/project/duli/content_relation_ana/utils')
import jsonlines
#from utils import *
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.colors import ListedColormap
from matplotlib.colors import ListedColormap, rgb2hex
import matplotlib.pyplot as plt
import random
def remove_adjacent_duplicates(lst):
    if not lst:
        return []
    result = [lst[0]]
    for num in lst[1:]:
        if num != result[-1]:
            result.append(num)
    return result

def generate_colors(n):
    """
    生成任意数量的颜色。
    
    参数:
    n (int): 需要生成的颜色数量。
    
    返回:
    list: 颜色列表，包含n个颜色。
    """
    colors = plt.cm.hsv(np.linspace(0, 1, n))
    return [rgb2hex(c) for c in colors]

def create_listed_colormap(n):
    """
    创建包含任意数量颜色的 ListedColormap。
    
    参数:
    n (int): 需要生成的颜色数量。
    
    返回:
    ListedColormap: 自定义的颜色映射。
    """
    colors = generate_colors(n)
    return ListedColormap(colors)


def distrib_ana_tsne(data_tsne, cate_ids, cate_names=None, out_path='./tsne.png', s=2, continue_draw=False):
    # assert len(cate_ids) == len(cate_names)
    if not cate_names:
        cate_names = remove_adjacent_duplicates(cate_ids)
    cmap1 = plt.get_cmap('tab20b')
    cmap2 = plt.get_cmap('tab20c')
    NUM_CATES = len(set(cate_ids))
    if len(cmap1.colors) + len(cmap2.colors) >= NUM_CATES:
        new_cmap = ListedColormap(random.sample((cmap1.colors+cmap2.colors), NUM_CATES)[:NUM_CATES])
    else:
        new_cmap = create_listed_colormap(NUM_CATES)
    plt.figure(figsize=(15,9), dpi=300)  # Make the plot flatter
    plt.subplots_adjust(left=0.1, right=0.7, top=0.9, bottom=0.2)
    scatter = plt.scatter(data_tsne[:, 0], data_tsne[:, 1], c=cate_ids, cmap=new_cmap, marker='o', s=s, alpha=0.9)
    cbar = plt.colorbar(scatter)
    cbar.set_ticks([i-float(i)/NUM_CATES for i in range(NUM_CATES)])
    cbar.set_ticklabels(cate_names)
    cbar.set_label('Cluster Label')
    plt.title('DBSCAN Clustering Over t-SNE')
    plt.xlabel('t-SNE Feature 1')
    plt.ylabel('t-SNE Feature 2')
    # plt.savefig('./tsne_llmsys_llmsysseed_1m_152k.png')
    if not continue_draw:
        plt.savefig(out_path)
